package srcBayes;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Random;
import java.util.Vector;

import lib.Pair;
import lib.Utility;
import lib.logging.ANSIConsoleLogger;
import lib.logging.ILogger;
import InferenceGraphs.InferenceGraph;
import InferenceGraphs.InferenceGraphNode;

import com.google.common.base.Stopwatch;


public class BayesInference
{
    private static final int BurnInIerations = 100000;
    private static final int MaxIterations = 10000000;
    private static final int SampleSkipCount = 100;
    private static final int SnapInterval = 1000000;
    private static boolean PrintSnapShots = false;

    public static void main(String[] args) throws Exception
    {
        ILogger logger = new ANSIConsoleLogger();

        try
        {
            if (args.length != 3)
            {
                logger.LogError(
                        "Incorrect number of paramets: %d. " +
                        "Usage: java BayesInference <bif file> <path to input parameters file> <output file>",
                        args.length);

                System.exit(1);
            }

            String inputBifFile = args[0];
            String inputFilePath = args[1];
            String outputFilePath = args[2];

            InferenceGraph inferenceGraph = new InferenceGraph(inputBifFile);
            BayesInference bayesInference = new BayesInference(logger, inputFilePath, inferenceGraph);

            Stopwatch stopwatch = new Stopwatch().start();

            bayesInference.ComputeInferences(outputFilePath);

            stopwatch.stop();
            long milis = stopwatch.elapsedMillis();

            System.out.println("Time elapsed: " + (milis / 1000.0));
        }
        catch (Exception ex)
        {
            logger.LogException(ex);
            throw ex;
        }
    }

    private static HashMap<String, String> ParseParametersFromInputFile(String filePath) throws IOException
    {
        HashMap<String, String> map = new HashMap<String, String>();

        BufferedReader br = new BufferedReader(new FileReader(filePath));
        String line = null;
        while ((line = br.readLine()) != null)
        {
            line = line.trim();

            if (line.equals("."))
            {
                break;
            }

            String[] splits = line.split("=");

            String parameter = splits[0].trim();
            parameter = parameter.replaceAll("\"", "");

            String value = splits[1].trim();
            value = value.replaceAll("\"", "");

            map.put(parameter, value);
        }

        return map;
    }

    private static void ValidateGivenParameters(
            HashMap<String, String> parameters,
            HashMap<String, InferenceGraphNode> nameToNodeMapping)
    {
        for (String parameter : parameters.keySet())
        {
            String givenValue = parameters.get(parameter);
            String[] possibleValues = nameToNodeMapping.get(parameter).get_values();
            boolean matchFound = false;

            for (String possibleValue : possibleValues)
            {
                if (possibleValue.equals(givenValue))
                {
                    matchFound = true;
                }
            }

            if (!matchFound)
            {
                throw new IllegalArgumentException(
                        String.format(
                                "Parameter '%s' has invalid value %s. Allowed values: %s",
                                parameter,
                                givenValue,
                                Utility.FlattenArrayToString(possibleValues)));
            }
        }
    }

    private HashMap<String, String> givenParameters;

    private HashMap<String, InferenceGraphNode> nameToNodeMapping;

    private List<String> diagnoses;

    private InferenceGraph inferenceGraph;

    private ILogger logger;

    private Random randomGenerator;

    public BayesInference(ILogger logger, String inputFilePath, InferenceGraph inferenceGraph) throws IOException
    {
        this.logger = logger;
        this.givenParameters = ParseParametersFromInputFile(inputFilePath);
        this.inferenceGraph = inferenceGraph;
        this.diagnoses = GetListOfDiagnoses();
        this.nameToNodeMapping = PopulateNameToNodeMapping(inferenceGraph);
        this.randomGenerator = new Random();

        ValidateGivenParameters(givenParameters, nameToNodeMapping);
    }

    private String ChooseValueAccordingToProbabilites(
            String[] possibleValues,
            double[] probabilities)
    {
        // probabilities[] contains the probability mass function (PMF),
        // for the possibleValues[] (http://en.wikipedia.org/wiki/Probability_mass_function)
        // We need to choose a value based on the PMF.
        // Algo described at: http://en.wikipedia.org/wiki/Pseudo-random_number_sampling#Finite_discrete_distributions

        int size = possibleValues.length;

        List<Pair<Double, Double>> intervals = new ArrayList<Pair<Double,Double>>();

        // compute cumulative distribution function (CDF) first
        double[] cdf = new double[size];
        cdf[0] = probabilities[0];
        intervals.add(new Pair<Double, Double>(0.0, cdf[0]));

        for (int i = 1; i < size; ++i)
        {
            cdf[i] = cdf[i - 1] + probabilities[i];
            intervals.add(new Pair<Double, Double>(cdf[i - 1], cdf[i]));
        }

        double randomValue = this.randomGenerator.nextDouble();

        for (int i = 0; i < intervals.size(); ++i)
        {
            Pair<Double, Double> interval = intervals.get(i);

            if (randomValue >= interval.A() && randomValue < interval.B())
            {
                return possibleValues[i];
            }
        }

        throw new IllegalStateException("Unable to choose a value for random variable!");
    }

    private void ComputeInferences(String outputFile) throws IOException
    {
        List<String> nonObservedVariables = GetNonObservedVariables();
        HashMap<String, String> nonObservedVariablesValues = GenerateNonObservedVariableValuesRandomly(nonObservedVariables);

        Sample currentSample = new Sample(this.givenParameters, nonObservedVariablesValues);

        // throw away initial samples during burn-in
        for (int i = 0; i < BurnInIerations; ++i)
        {
            UpdateSample(currentSample, nonObservedVariables);
        }

        // now start capturing samples
        int sampleCount = 0;
        HashMap<String, Integer> frequencyTable = InitializeFrequencyTable();

        for (int i = 0; i < MaxIterations; ++i)
        {
            UpdateSample(currentSample, nonObservedVariables);

            if ((i % SampleSkipCount) == 0)
            {
                UpdateFrequencyTable(frequencyTable, currentSample);
                ++sampleCount;
            }

            if (PrintSnapShots && (i % SnapInterval) == 0)
            {
                logger.LogInfo("Snap after iteration: %d, SampleCount: %d", i, sampleCount);
                PrintSnapShot(frequencyTable, sampleCount);
            }
        }

        OutputInferencesToFile(outputFile, frequencyTable, sampleCount);
    }

    private String GenerateFrequencyTableKey(String name, String value)
    {
        return (name + "_" + value);
    }

    private HashMap<String, String> GenerateNonObservedVariableValuesRandomly(
            List<String> nonObservedVariables)
    {
        HashMap<String, String> map = new HashMap<String, String>();

        for (String varName : nonObservedVariables)
        {
            InferenceGraphNode node = GetInferenceNode(varName);
            String[] values = node.get_values();
            int randomIndex = this.randomGenerator.nextInt(values.length);
            map.put(varName, values[randomIndex]);
        }

        return map;
    }

    private double GetConditionalProbability(
            InferenceGraphNode node,
            String value,
            Sample currentSample)
    {
        // instead of doing product of probabilities,
        // do sum of log of probabilities, and then return log inverse.
        // this way we can avoid underflow to 0 on multiplying small probabilities.

        double sum = Math.log10(LookupProbabilityForNode(node, value, currentSample));

        Vector children = node.get_children();
        for (int i = 0; i < children.size(); ++i)
        {
            InferenceGraphNode childNode = (InferenceGraphNode) children.elementAt(i);
            String childValue = currentSample.GetVariableValue(childNode.get_name());
            sum += Math.log10(LookupProbabilityForChild(node, value, childNode, childValue, currentSample));
        }

        return Math.pow(10, sum);
    }

    private InferenceGraphNode GetInferenceNode(String name)
    {
        if (!this.nameToNodeMapping.containsKey(name))
        {
            throw new IllegalArgumentException(
                    String.format(
                            "Variable '%s' does not belong to inference graph!",
                            name));
        }

        return this.nameToNodeMapping.get(name);
    }

    private List<String> GetListOfDiagnoses()
    {
        List<String> diagnoses = new ArrayList<String>();

        diagnoses.add("Hypovolemia");
        diagnoses.add("LVFailure");
        diagnoses.add("Anaphylaxis");
        diagnoses.add("InsuffAnesth");
        diagnoses.add("PulmEmbolus");
        diagnoses.add("Intubation");
        diagnoses.add("KinkedTube");
        diagnoses.add("Disconnect");

        return diagnoses;
    }

    private List<String> GetNonObservedVariables()
    {
        List<String> vars = new ArrayList<String>();

        for (String varName : this.nameToNodeMapping.keySet())
        {
            if (!this.givenParameters.containsKey(varName))
            {
                vars.add(varName);
            }
        }

        return vars;
    }

    private String GetOutputLineString(String diagnosis, String[] values, double[] probabilities)
    {
        StringBuilder sb = new StringBuilder();

        sb.append("\"");
        sb.append(diagnosis);
        sb.append("\" = { ");

        for (int i = 0; i < values.length; ++i)
        {
            sb.append(values[i]);
            sb.append(" : ");
            sb.append(probabilities[i]);
            sb.append(" ");
        }

        sb.append("}");

        return sb.toString();
    }

    private HashMap<String, Integer> InitializeFrequencyTable()
    {
        HashMap<String, Integer> map = new HashMap<String, Integer>();

        for (String diagnosis : this.diagnoses)
        {
            InferenceGraphNode node = GetInferenceNode(diagnosis);
            String[] values = node.get_values();

            for (String value : values)
            {
                String key = GenerateFrequencyTableKey(diagnosis, value);
                map.put(key, 0);
            }
        }

        return map;
    }

    private double LookupProbabilityForChild(
            InferenceGraphNode parentNode,
            String parentValue,
            InferenceGraphNode childNode,
            String childValue,
            Sample currentSample)
    {
        Vector parents = childNode.get_parents();

        int arraySize = parents.size() + 1;
        String[][] s = new String[arraySize][2];

        s[0][0] = childNode.get_name();
        s[0][1] = childValue;

        for (int i = 1; i < arraySize; ++i)
        {
            InferenceGraphNode node = (InferenceGraphNode) parents.elementAt(i - 1);
            s[i][0] = node.get_name();

            if (node.get_name().equals(parentNode.get_name()))
            {
                s[i][1] = parentValue;
            }
            else
            {
                s[i][1] = currentSample.GetVariableValue(node.get_name());
            }
        }

        return childNode.get_function_value(s);
    }

    private double LookupProbabilityForNode(
            InferenceGraphNode node,
            String value,
            Sample currentSample)
    {
        Vector parents = node.get_parents();

        int arraySize = parents.size() + 1;
        String[][] s = new String[arraySize][2];

        s[0][0] = node.get_name();
        s[0][1] = value;

        for (int i = 1; i < arraySize; ++i)
        {
            InferenceGraphNode parentNode = (InferenceGraphNode) parents.elementAt(i - 1);
            s[i][0] = parentNode.get_name();
            s[i][1] = currentSample.GetVariableValue(parentNode.get_name());
        }

        return node.get_function_value(s);
    }

    private double[] NormalizeProbabilities(double[] probabilities)
    {
        double[] normalizedProbabilities = new double[probabilities.length];

        double sum = 0.0;

        for (double probability : probabilities)
        {
            sum += probability;
        }

        for (int i = 0; i < probabilities.length; ++i)
        {
            normalizedProbabilities[i] = probabilities[i] / sum;
        }

        return normalizedProbabilities;
    }

    private void OutputInferencesToFile(
            String outputFile,
            HashMap<String, Integer> frequencyTable,
            int sampleCount) throws IOException
    {
        BufferedWriter bufWriter = new BufferedWriter(
                new FileWriter(outputFile));

        for (String diagnosis : this.diagnoses)
        {
            InferenceGraphNode node = GetInferenceNode(diagnosis);
            String[] values = node.get_values();
            double[] probabilities = new double[values.length];

            for (int i = 0; i < values.length; ++i)
            {
                String key = GenerateFrequencyTableKey(diagnosis, values[i]);
                probabilities[i] = ((double) frequencyTable.get(key)) / ((double) sampleCount);
            }

            bufWriter.write(GetOutputLineString(diagnosis, values, probabilities));
            bufWriter.newLine();
        }

        bufWriter.close();

    }

    private String PickValue(String variableToFlip, Sample currentSample)
    {
        InferenceGraphNode node = GetInferenceNode(variableToFlip);
        String[] possibleValues = node.get_values();
        double[] probabilities = new double[possibleValues.length];

        for (int i = 0; i < possibleValues.length; ++i)
        {
            probabilities[i] = GetConditionalProbability(node, possibleValues[i], currentSample);
        }

        double[] normalizedProbabilities = NormalizeProbabilities(probabilities);
        return ChooseValueAccordingToProbabilites(possibleValues, normalizedProbabilities);
    }

    private HashMap<String, InferenceGraphNode> PopulateNameToNodeMapping(
            InferenceGraph inferenceGraph)
    {
        HashMap<String, InferenceGraphNode> map = new HashMap<String, InferenceGraphNode>();

        Object[] nodes = this.inferenceGraph.get_nodes().toArray();
        for (Object node : nodes)
        {
            InferenceGraphNode inferenceNode = (InferenceGraphNode) node;
            String nodeName = inferenceNode.get_name();

            map.put(nodeName, inferenceNode);
        }

        return map;
    }

    private void PrintSnapShot(
            HashMap<String, Integer> frequencyTable,
            int sampleCount)
    {
        for (String diagnosis : this.diagnoses)
        {
            InferenceGraphNode node = GetInferenceNode(diagnosis);
            String[] values = node.get_values();
            double[] probabilities = new double[values.length];

            for (int i = 0; i < values.length; ++i)
            {
                String key = GenerateFrequencyTableKey(diagnosis, values[i]);
                probabilities[i] = ((double) frequencyTable.get(key)) / ((double) sampleCount);
            }

            logger.LogInfo(GetOutputLineString(diagnosis, values, probabilities));
        }
    }

    private void UpdateFrequencyTable(
            HashMap<String, Integer> frequencyTable,
            Sample sample)
    {
        for (String diagnosis : this.diagnoses)
        {
            String valueInSample = sample.GetVariableValue(diagnosis);
            String key = GenerateFrequencyTableKey(diagnosis, valueInSample);

            int currentCount = frequencyTable.get(key);
            frequencyTable.put(key, currentCount + 1);
        }
    }

    private void UpdateSample(
            Sample currentSample,
            List<String> nonObservedVariables)
    {
        int randomIndex = this.randomGenerator.nextInt(nonObservedVariables.size());

        String variableToFlip = nonObservedVariables.get(randomIndex);
        String newValue = PickValue(variableToFlip, currentSample);

        currentSample.SetNonObservedVariable(variableToFlip, newValue);
    }

}
